"""
根据SQLite中已有【标题】和文档更新ChromaDB
"""

from db_utils import Database, VectorStore, Document


# 创建数据库连接和向量存储实例
db = Database()
session = db.connect()
vector_store = VectorStore(db)


def update_all_document_embeddings():
    print("重新计算所有文档的嵌入向量，包含标题信息...")
    print("============================")

    # 获取所有文档
    documents = session.query(Document).all()

    print(f"找到 {len(documents)} 个文档，开始更新嵌入向量...")

    for i, document in enumerate(documents, 1):
        # 获取文档的分类ID列表
        category_ids = [cat.id for cat in document.categories]

        # 将标题和内容合并后进行嵌入
        text_to_embed = document.content
        if document.title:
            text_to_embed = f"{document.title}\n{document.content}"

        # 计算新的嵌入向量
        embedding = vector_store.get_text_embedding(text_to_embed)

        # 更新ChromaDB中的嵌入向量和元数据
        metadata = {
            "document_id": document.id,
            "category_ids": ",".join(map(str, category_ids)) if category_ids else "",
            "title": document.title if document.title else "",
            "content": document.content,
        }

        # 先删除旧的向量，再添加新的向量
        try:
            vector_store.collection.delete(ids=[str(document.id)])
        except Exception as e:
            print(f"  删除文档 {document.id} 的旧向量时出错: {e}")

        vector_store.collection.add(
            ids=[str(document.id)], embeddings=[embedding], metadatas=[metadata]
        )

        print(
            f"  已更新文档 {i}/{len(documents)}: ID={document.id}, 标题='{document.title or '无标题'}'"
        )

    print("\n所有文档的嵌入向量已更新完成！")
    print("现在搜索功能将同时基于标题和内容进行检索。")


# 测试更新后的搜索功能
def test_search():
    print("\n测试更新后的搜索功能:")
    print("====================")

    # 测试标题搜索
    query = "西游记"
    print(f"\n搜索关键词: '{query}' (应该优先匹配标题包含'西游记'的文档)")
    results = vector_store.search_similar_documents(query, top_k=5)

    for i, (doc_id, content, title, similarity, category_ids) in enumerate(results):
        print(f"  {i+1}. ID: {doc_id}, 标题: '{title}', 相似度: {similarity:.4f}")


if __name__ == "__main__":
    try:
        # 运行更新脚本
        update_all_document_embeddings()
        # 执行测试
        test_search()
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        db.disconnect(session)
